package com.example.demo.service;

import com.example.demo.util.JsonUtil;
import com.google.common.collect.Lists;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.vector.request.AnnSearchReq;
import io.milvus.v2.service.vector.request.HybridSearchReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.request.ranker.WeightedRanker;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;

@Slf4j
@Service
public class MilvusService extends AbstractMilvusService {

    @Autowired
    private MilvusClientV2 client;

    @Autowired
    private EmbedingService embedingService;

    @Override
    public void createCollection() {
        CreateCollectionReq createCollectionReq  = CreateCollectionReq.builder()
                .collectionName(CollectionName)
                .collectionSchema(schema())
                .indexParams(index())
                .dimension(1024).build();
        client.createCollection(createCollectionReq);
    }

    @Override
    public CreateCollectionReq.CollectionSchema schema() {
        CreateCollectionReq.CollectionSchema schema = client.createSchema();

        // 3.2 Add fields to schema
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("content")
                .dataType(DataType.VarChar)
                .maxLength(2000)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("content_vector")
                .dataType(DataType.FloatVector)
                .dimension(1024)
                .build());

        log.info ("schema创建完成>>>> schema:{}", JsonUtil.toJsonString(schema));
        return schema;
    }

    @Override
    public List<IndexParam> index() {
        log.info("文章索引创建完成>>>> index:{}", JsonUtil.toJsonString(IndexParam.builder().fieldName("content_vector").metricType(IndexParam.MetricType.COSINE).build()));
        return List.of(IndexParam.builder().fieldName("content_vector").metricType(IndexParam.MetricType.COSINE).build());
    }

    @Override
    public void insert(String contnt) throws Exception {
        // 批量embedding
        List<float[]> contentVectors = embedingService.embedding(contnt);
        List<JsonObject> vectors = new ArrayList<>();
        for (float[] vector : contentVectors) {
            JsonObject object = new JsonObject();
            object.addProperty("content", contnt);

            Gson gson = new Gson();
            object.add("content_vector", gson.toJsonTree(vector));

            vectors.add(object);
        }

        InsertReq request = InsertReq.builder()
                .collectionName(CollectionName)
                .data(vectors)
                .build();
        InsertResp resp = client.insert(request);

        log.info("文章向量数据入库完成>>>> resp:{}", JsonUtil.toJsonString(resp));
    }

    @Override
    public List<SearchResp.SearchResult> search(String contnt) throws Exception {
        List<float[]> contentEmbed = embedingService.embedding(contnt);

        List<FloatVec> floatVecs = new ArrayList<>();
        for  (float[] vector : contentEmbed) {
            floatVecs.add(new FloatVec(vector));
        }

        AnnSearchReq annSearchReq = AnnSearchReq.builder().metricType(IndexParam.MetricType.COSINE)
                                                          .vectors(Collections.unmodifiableList(floatVecs))
                                                          .topK(1).vectorFieldName("content_vector").build();

        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(CollectionName)
                .searchRequests(Lists.newArrayList(annSearchReq))
                .ranker(new WeightedRanker(Arrays.asList(0.5f)))
                .outFields(Lists.newArrayList("content"))
                .topK(1) // 返回 top K 结果
                .build();

        SearchResp searchResp = client.hybridSearch(hybridSearchReq);

        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();


        return searchResults.get(0);
    }
}
